#include "hzpch.h"
#include "Renderer.h"

#include "Renderer2D.h"
#include "Platform/OpenGL/OpenGLShader.h"
#include "Hazel/Core/Random.h"

#include <execution>

namespace Hazel
{

	Renderer::SceneData* Renderer::s_SceneData = new Renderer::SceneData;

	void Renderer::Init()
	{
		RenderCommand::Init();
		Renderer2D::Init();
	}

	void Renderer::OnWindowResize(uint32_t width, uint32_t height)
	{
		RenderCommand::SetViewport(0, 0, width, height);
	}

	void Renderer::OnResize(uint32_t width, uint32_t height)
	{
		if (m_FinalImage)
		{
			if (m_FinalImage->GetHeight() == height && m_FinalImage->GetWidth() == width)
				return;
			m_FinalImage->ReSize(width, height);
		}
		else
		{
			m_FinalImage = Texture2D::Create(width, height);
		}
		delete[] m_imageData;
		m_imageData = new uint32_t[width * height];		
		
		delete[] m_AccumulationData;
		m_AccumulationData = new glm::vec4[width * height];

		ResetFrameIndex();

		m_ImageHorizontalIter.resize(width);
		m_ImageVerticalIter.resize(height);

		for (uint32_t i = 0; i < width; i++)
			m_ImageHorizontalIter[i] = i;
		for (uint32_t i = 0; i < height; i++)
			m_ImageVerticalIter[i] = i;
	}

	void Renderer::Render(const Scene& scene, const Camera& camera)
	{
		m_ActivateCamera = &camera;
		m_ActivateScene = &scene;
		if (m_FrameIndex == 1)
			memset(m_AccumulationData, 0, m_FinalImage->GetWidth() *m_FinalImage->GetHeight() * sizeof(glm::vec4));
#define MT 1
#if MT

		std::for_each(std::execution::par,m_ImageVerticalIter.begin(), m_ImageVerticalIter.end(),
			[&](uint32_t y)
			{		
				std::for_each(m_ImageHorizontalIter.begin(), m_ImageHorizontalIter.end(),
					[this, y](uint32_t x)
					{
						glm::vec4 color = PerPixel(x, y);

						m_AccumulationData[x + y * m_FinalImage->GetWidth()] += color;

						glm::vec4 accumulatedColor = m_AccumulationData[x + y * m_FinalImage->GetWidth()];
						accumulatedColor /= (float)m_FrameIndex;

						accumulatedColor = glm::clamp(accumulatedColor, glm::vec4(0.0f), glm::vec4(1.0f));
						m_imageData[x + y * m_FinalImage->GetWidth()] = ConvertToRGBA(accumulatedColor);
					}); 

			});
#else

		for (uint32_t y = 0; y < m_FinalImage->GetHeight(); y++)
			for (uint32_t x = 0; x < m_FinalImage->GetWidth(); x++)
			{

				glm::vec4 color = PerPixel(x, y);

				m_AccumulationData[x + y * m_FinalImage->GetWidth()] += color;

				glm::vec4 accumulatedColor = m_AccumulationData[x + y * m_FinalImage->GetWidth()];
				accumulatedColor /= (float)m_FrameIndex;

				accumulatedColor = glm::clamp(accumulatedColor, glm::vec4(0.0f), glm::vec4(1.0f));
				m_imageData[x + y * m_FinalImage->GetWidth()] = ConvertToRGBA(accumulatedColor);
			}
#endif // MT
		m_FinalImage->SetData(m_imageData, m_FinalImage->GetHeight() * m_FinalImage->GetWidth() * 4);

		if (m_Settings.Accumulate)
		{
			m_FrameIndex++;
		}
		else
		{
			m_FrameIndex = 1;
		}

	}

	glm::vec4 Renderer::PerPixel(uint32_t x, uint32_t y)
	{
		Ray ray;
		ray.Origin = m_ActivateCamera->GetPosition();
		ray.Direction = m_ActivateCamera->GetRayDirections()[x + y * m_FinalImage->GetWidth()];
		 
		glm::vec3 color(0.0f);
		float multiplier = 1.0f;
		int bounces = 5;
		for (int i = 0; i < bounces; i++)
		{
			Renderer::HitPayload payload = TraceRay(ray);
			if (payload.HitDistance < 0.0f)
			{
				glm::vec3 skycolor = glm::vec3(0.6f, 0.7f, 0.9f);
				color += skycolor * multiplier;
				break;
			}

			glm::vec3 lightDir = glm::normalize(glm::vec3(-1, -1, -1));
			float d = glm::max(glm::dot(payload.WorldNormal, -lightDir), 0.0f);

			const Sphere& sphere = m_ActivateScene->Spheres[payload.ObjectIndex];
			const Material& material = m_ActivateScene->Materials[sphere.MaterialIndex];
			glm::vec3 sphereColor = material.Albedo;
			sphereColor *= d;
			color += multiplier * sphereColor;
			multiplier *= 0.5f;

			ray.Origin = payload.WorldPosition + payload.WorldNormal * 0.0001f;
			ray.Direction = glm::reflect(ray.Direction, payload.WorldNormal + material.Roughness * Random::Vec3(-0.5f,0.5f));
		}
		return glm::vec4(color, 1);
	}

	Renderer::HitPayload Renderer::TraceRay(const Ray& ray)
	{

		int closestSphere = -1;
		//float hitDistance = FLT_MAX;
		float hitDistance = std::numeric_limits<float>::max();
		for (size_t i = 0; i < m_ActivateScene->Spheres.size(); i++)
		{
			const Sphere& sphere = m_ActivateScene->Spheres[i];
			glm::vec3 origin = ray.Origin - sphere.Position;

			float a = glm::dot(ray.Direction, ray.Direction);
			//float a = rayDirection.x * rayDirection.x + rayDirection.y * rayDirection.y + rayDirection.z * rayDirection.z;
			float b = 2.0f * glm::dot(origin, ray.Direction);
			float c = glm::dot(origin, origin) - sphere.Radius * sphere.Radius;

			float discriminant = b * b - 4.0f * a * c;
			if (discriminant < 0.0f)
				continue;

			float closestT = (-b - glm::sqrt(discriminant)) / (2.0f * a);
			//float t1 = (-b + glm::sqrt(discriminant)) / (2.0f * a);
			if ( closestT > 0.0f && closestT < hitDistance)
			{
				hitDistance = closestT;
				closestSphere = (int)i;
			}
		}

		if (closestSphere < 0)
		{
			return Miss(ray);
		}
		return ClosestHit(ray, hitDistance, closestSphere);

	}

	Renderer::HitPayload Renderer::ClosestHit(const Ray& ray, float hitDistance, int objectIndex)
	{
		Renderer::HitPayload payload;
		payload.HitDistance = hitDistance;
		payload.ObjectIndex = objectIndex;

		const Sphere& closestSphere = m_ActivateScene->Spheres[objectIndex];

		glm::vec3 origin = ray.Origin - closestSphere.Position;

		payload.WorldPosition = origin + ray.Direction * hitDistance;
		payload.WorldNormal = glm::normalize(payload.WorldPosition);

		payload.WorldPosition += closestSphere.Position;

		return payload;
	}

	Renderer::HitPayload Renderer::Miss(const Ray& ray)
	{
		Renderer::HitPayload payload;
		payload.HitDistance = -1.0f;
		return payload;
	}

	uint32_t Renderer::ConvertToRGBA(glm::vec4 color)
	{
		uint8_t r = (uint8_t)(color.r * 255.0f);
		uint8_t g = (uint8_t)(color.g * 255.0f);
		uint8_t b = (uint8_t)(color.b * 255.0f);
		uint8_t a = (uint8_t)(color.a * 255.0f);

		return (a << 24) | (b << 16) | (g << 8) | r;
	}

	void Renderer::BeginScene(OrthographicCamera& camera)
	{
		s_SceneData->ViewProjectionMatrix = camera.GetVPMatrix();
	}

	void Renderer::EndScene()
	{
	}

	void Renderer::Submit(const Ref<VertexArray>& vertexArray, const Ref<Shader>& shader, const glm::mat4& transform)
	{
		shader->Bind();
		std::dynamic_pointer_cast<OpenGLShader> (shader)->UploadUniformMat4("u_ViewProjection", s_SceneData->ViewProjectionMatrix);
		std::dynamic_pointer_cast<OpenGLShader> (shader)->UploadUniformMat4("u_Transform", transform);


		vertexArray->Bind();
		RenderCommand::DrawIndexed(vertexArray);
	}

}
